#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © 2020 wanghch <wanghch@wanghch-pc>
#
# Distributed under terms of the MIT license.

"""

"""
import baostock as bs
import pandas as pd
import numpy as np
import argparse
import datetime
from os.path import join
import os
import sys


data_dir = 'data'


today_t = datetime.datetime.today()
today = today_t.strftime('%Y-%m-%d')

class BaseStockApi(object):
    k_fields = "date,code,open,high,low,close,volume,amount,adjustflag,turn,tradestatus,pctChg,peTTM,pbMRQ,psTTM,pcfNcfTTM,isST"
    def __init__(self):
        bs.login()
        self.last_trade_day = self.get_last_trade_day()

    def get_last_trade_day(self):
        day30ago = (today_t - datetime.timedelta(days=30)).strftime("%Y-%m-%d")
        yesterday = today_t.strftime("%Y-%m-%d")
        rs = bs.query_trade_dates(start_date=day30ago, end_date = yesterday)
        date = [v.calendar_date for i, v in rs.get_data().iterrows() if v.is_trading_day == '1']

        data = bs.query_all_stock(date[-1]).get_data()
        if data.shape[0] == 0:
            return date[-2]
        else:
            return date[-1]


    # get all stock info
    def get_all_stocks(self):
        return bs.query_all_stock(self.last_trade_day).get_data()

    def save_to_csv(self, result, csv_file, encoding = 'utf-8'):
        result.to_csv(csv_file, encoding = encoding, index=False)

    def save_all_stock(self):
        rt = api.get_all_stocks()
        self.save_to_csv(rt, 'all_stock.csv')

    def merge(self, code = None, merge_mode = "single", merge_dir = "merge"):
        code_files = None
        if code and hasattr(self, code):
            code_df = getattr(self, code)()
            codes = code_df.code.values
            codes_files = (code_df.code + "." + code_df.code_name+".csv").values
        if code_files:
            merge_mode = "split"


        print(code_files)
        all_base_dir = join(data_dir, 'history_k', 'all')
        file_dirs = os.listdir(all_base_dir)
        file_dirs.sort(reverse=True)

        new_dir = join(all_base_dir, file_dirs[0])

        if code:
            merge_base_dir = join(data_dir, 'history_k', merge_dir, code)
        else:
            merge_base_dir = join(data_dir, 'history_k', merge_dir, file_dirs[0].split("_")[1])
        if os.path.exists(merge_base_dir):
            print("merge base dir %s exists" % merge_base_dir)
            return
        else:
            os.makedirs(merge_base_dir)

        if code:
            file_names = codes_files
        else:
            file_names = os.listdir(new_dir)


        if code and merge_mode == 'single':
            file_dirs.sort(reverse = False)
            csv_dfs = []
            for fn in file_names:
                print(fn)
                for fd in file_dirs:
                    csv_f = join(all_base_dir, fd, fn)
                    if os.path.exists(csv_f):
                        csv_dfs.append(pd.read_csv(csv_f))
            merge_df = pd.concat(csv_dfs, ignore_index = True)
            merge_df.to_csv(join(merge_base_dir, code + ".csv"), index=False)
        else:
            file_dirs.sort(reverse = False)
            for fn in file_names:
                print(fn)
                csv_dfs = []
                for fd in file_dirs:
                    csv_f = join(all_base_dir, fd, fn)
                    if os.path.exists(csv_f):
                        csv_dfs.append(pd.read_csv(csv_f))
                merge_df = pd.concat(csv_dfs, ignore_index = True)
                merge_df.to_csv(join(merge_base_dir, fn), index=False)




    def get_date_diff(self, d1, d2):
        day1 = datetime.datetime.strptime(d1, '%Y-%m-%d')
        day2 = datetime.datetime.strptime(d2, '%Y-%m-%d')
        return (day1 - day2).days


    def get_all(self, start_date, end_date = None):
        if end_date is None:
            end_date = self.last_trade_day

        if start_date is None:
            all_base_dir = join(data_dir, 'history_k', 'all')
            files = os.listdir(all_base_dir)
            x = max([fn.split("_")[1] for fn in files if "_" in fn])
            last_add1_day = datetime.datetime.strptime(x, '%Y-%m-%d') + datetime.timedelta(days = 1)
            start_date = last_add1_day.strftime("%Y-%m-%d")

        if self.get_date_diff(start_date, end_date) >= 0:
            print("start date: %s >= end date: %s, skip" % (start_date , end_date))
            return

        if start_date == end_date:
            print("start date == end date: %s, skip" % end_date)
            return
        all_dir = join(data_dir, 'history_k', 'all', start_date+ '_' +end_date)
        if not os.path.exists(all_dir):
            os.makedirs(all_dir)

        # all_stocks = self.get_all_stocks()
        all_stocks = self.hs300()
        # print(all_stocks)
        for index, row in all_stocks.iterrows():
            code = row['code']
            print(code)
            code_name = row['code_name']
            print(code , code_name)
            self.get_history_k(code, start_date, end_date, join(all_dir, code+'.'+code_name+'.csv'))


    def get_history_k(self, code, start_date, end_date = None, file_path = None):
        if end_date is None:
            end_date = self.last_trade_day
        if os.path.exists(file_path):
            return pd.read_csv(file_path)
        else:
            rs = bs.query_history_k_data_plus(code,
               self.k_fields,
                    start_date=start_date, end_date=end_date,
            frequency="d", adjustflag = "2")
            result = rs.get_data()
            result.to_csv(file_path, index=False)
            return result

    def get_industry(self):
        df = bs.query_stock_industry().get_data()
        file_path = join(data_dir, 'industry.csv')
        df.to_csv(file_path, index=False)


    def sz50(self):
        file_path = join(data_dir, 'sz50.csv')
        if os.path.exists(file_path):
            return pd.read_csv(file_path)
        else:
            df = bs.query_sz50_stocks().get_data()
            df.to_csv(file_path, index=False)
            return df

    def hs300(self):
        file_path = join(data_dir, 'hs300.csv')
        print(file_path)
        if os.path.exists(file_path):
            return pd.read_csv(file_path)
        else:
            df = bs.query_hs300_stocks().get_data()
            df.to_csv(file_path, index=False)
            return df

    def zz500(self):
        file_path = join(data_dir, 'zz500.csv')
        if os.path.exists(file_path):
            return pd.read_csv(file_path)
        else:
            df = bs.query_zz500_stocks().get_data()
            df.to_csv(file_path, index=False)
            return df


    def logout(self):
        bs.logout()


def get_all_args():
    parser.add_argument('--start_date', type=str, help='start date')
    parser.add_argument('--end_date', type=str, help='end date')

def get_history_k_args():
    parser.add_argument('--code', type=str, help='code')
    parser.add_argument('--start_date', type=str,  help='start date')
    parser.add_argument('--end_date', type=str, help='end date')
    parser.add_argument('--file_path', type=str, help='file_path')

def merge_args():
    parser.add_argument('--code', type=str, help='code')
    parser.add_argument('--merge_mode', type=str, default='single', help='merge mode, single or split')
    parser.add_argument('--merge_dir', type=str, default='merge', help='merge')




if __name__ == '__main__':
    if len(sys.argv) == 1:
        print("need job")
        sys.exit(-1)
    job = sys.argv[1]
    cur_mod = sys.modules[__name__]

    parser = argparse.ArgumentParser(description= job+' stock api tools')
    if hasattr(cur_mod, job + "_args"):
        getattr(cur_mod, job + "_args")()

    job_args = sys.argv[2:]

    args = parser.parse_args(job_args)
    if not os.path.isdir(data_dir):
        os.makedirs(data_dir)


    api = BaseStockApi()
    kwargs = args.__dict__
    if hasattr(api, job):
        getattr(api, job)(**kwargs)

    api.logout()
